package org.emailclassifier.classifiers;

import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;

import org.emailclassifier.preprocessing.DataSetDocument;

public class AIMClassifier extends ClassifierBase {

	/**
	 * Category indexes for the probability matrix
	 */
	HashMap<String, Integer> m_catIndex = null;
	HashMap<Integer, String> m_catIndexi = null;
	HashMap<String,HashMap<String, Double>> m_vectorWPerCategory = null;
	
	/**
	 * Term indexes for the probability matrix
	 */
	HashMap<String, Integer> m_termIndex = null;
	
	/**
	 * Vocabulary used by the classifier.
	 * The amount of words used can be limited using setCutoff
	 * 
	 * @see setCutoff
	 */
	Set<String> m_vocabulary = null;
	
////////////////////////////////////////////////////////////////////

	public void train() {
		
		m_vectorWPerCategory = new HashMap<String,HashMap<String, Double>>();
		
		
		
		for( String cat : m_trainingSet.getCategories() ) {
			
			// iterate over every term in a category
			// compute category vectors
			
			HashMap<String,Double> vectorW = new HashMap<String,Double>();
			
			for( String term : m_trainingSet.getCategoryVocabulary(cat) ) {
				double tf = m_trainingSet.getCategoryTermFrequency( cat, term);
				double idf = 1 / Math.pow( m_trainingSet.getDocumentFrequency(term), 2 );
				vectorW.put(term, tf * idf );
			}
			m_vectorWPerCategory.put(cat, vectorW);
			
		}
		m_isTrained = true;
	}

	public String[] determineClasses(DataSetDocument content) {
		
		m_catIndexi = new HashMap<Integer, String>();
		
		if (!m_isTrained)
		{
			return null;
		}
		
		int catCount = m_trainingSet.getCategoryCount();
		
		double scores[] = new double[catCount];
		Map<String, Integer> documentWf = content.getTermFrequencies();
		
		int ii = 0;
		for (String category : m_trainingSet.getCategories()) {
			m_catIndexi.put(ii, category);
			scores[ii] = computeSim4( documentWf, category );
			ii++;
		}
		
		// DEBUG
		//System.out.printf("%f - %f - %f - %f - %f\n", scores[0], scores[1], scores[2], scores[3], scores[4]);
		
		HashMap<String, Double> unsorted = new HashMap<String, Double>();
		
		for (int i = 0; i<catCount; i++)
		{
			unsorted.put(m_catIndexi.get(i), scores[i]);
		}
		
		ClassProbComparator comp = new ClassProbComparator(unsorted);
		@SuppressWarnings("unchecked")
		TreeMap<String, Double> sortedCategories = new TreeMap<String, Double>(comp);
		
		sortedCategories.putAll(unsorted);
		
		return sortedCategories.keySet().toArray(new String[0]);
	}

	/**
	 * 
	 * @param vectorF Current document word frequency (wf)
	 * @return SIM4 score
	 */
	private double computeSim4(Map<String, Integer> vectorF, String category) {

		if( !m_vectorWPerCategory.containsKey(category) ) {
			return 0;
		}
		
		HashMap<String, Double> vectorW = m_vectorWPerCategory.get(category);
		double dotProductFW = 0;


		// iterate over every term in document
		// compute dot product of vectorF * vectorW
		for( String term : vectorF.keySet() ) {
			double freqW = vectorW.get(term) == null ? 0 : vectorW.get(term);
			dotProductFW += vectorF.get(term) * freqW;
		}

		
		return dotProductFW / Math.min( vectorF.size(), vectorW.size() );
	}

	public String getDescription() {
		return "AIM Classifier";
	}
}